#!/bin/bash

set -e

# 性能测试脚本 - 用于测试智能快照器的性能指标

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 默认配置
SNAPSHOTTER_NAME="smart"
TEST_IMAGES=(
    "docker.io/library/alpine:latest"
    "docker.io/library/nginx:latest"
    "docker.io/library/redis:latest"
    "docker.io/library/mysql:latest"
    "docker.io/library/postgres:latest"
)
ITERATIONS=5
WARMUP_ITERATIONS=2

# 日志和结果目录
RESULTS_DIR="./test-results"
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")
RESULTS_FILE="$RESULTS_DIR/performance-results-$TIMESTAMP.json"

# 创建结果目录
mkdir -p "$RESULTS_DIR"

# 记录日志函数
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $1" | tee -a "$RESULTS_DIR/performance-$TIMESTAMP.log"
}

error() {
    echo -e "${RED}[ERROR] $1${NC}" | tee -a "$RESULTS_DIR/performance-$TIMESTAMP.log"
}

success() {
    echo -e "${GREEN}[SUCCESS] $1${NC}" | tee -a "$RESULTS_DIR/performance-$TIMESTAMP.log"
}

info() {
    echo -e "${BLUE}[INFO] $1${NC}" | tee -a "$RESULTS_DIR/performance-$TIMESTAMP.log"
}

warning() {
    echo -e "${YELLOW}[WARNING] $1${NC}" | tee -a "$RESULTS_DIR/performance-$TIMESTAMP.log"
}

# 检查依赖
check_dependencies() {
    log "检查依赖项..."

    local deps=("ctr" "jq" "bc" "time")
    for dep in "${deps[@]}"; do
        if ! command -v "$dep" &> /dev/null; then
            error "$dep 命令未找到"
            exit 1
        fi
    done

    success "依赖项检查通过"
}

# 系统信息收集
collect_system_info() {
    info "收集系统信息..."

    local cpu_info=$(lscpu | grep "Model name" | cut -d: -f2 | xargs)
    local mem_info=$(free -h | grep "Mem:" | awk '{print $2}')
    local disk_info=$(df -h / | awk 'NR==2 {print $4}')

    cat > "$RESULTS_DIR/system-info-$TIMESTAMP.json" << EOF
{
    "timestamp": "$(date -Iseconds)",
    "cpu": "$cpu_info",
    "memory": "$mem_info",
    "disk_free": "$disk_info",
    "kernel": "$(uname -r)",
    "containerd_version": "$(ctr version | grep Version | head -1 | awk '{print $2}')"
}
EOF
}

# 拉取测试镜像
pull_test_images() {
    log "拉取测试镜像..."

    for image in "${TEST_IMAGES[@]}"; do
        if ! ctr image ls | grep -q "$(echo $image | cut -d: -f1 | cut -d/ -f3)"; then
            info "拉取镜像: $image"
            if ! timeout 600 ctr image pull "$image"; then
                warning "镜像拉取失败: $image"
            fi
        fi
    done
}

# 测量容器启动时间
measure_startup_time() {
    local image=$1
    local iteration=$2
    local container_name="perf-test-${iteration}"

    # 获取镜像的简单名称
    local image_name=$(echo "$image" | cut -d: -f1 | cut -d/ -f3)

    # 清理之前的容器
    ctr container rm "$container_name" > /dev/null 2>&1 || true

    # 测量启动时间
    local start_time=$(date +%s.%N)

    if timeout 60 ctr run --snapshotter "$SNAPSHOTTER_NAME" --rm "$image" "$container_name" echo "startup test" > /dev/null 2>&1; then
        local end_time=$(date +%s.%N)
        local duration=$(echo "$end_time - $start_time" | bc -l)

        echo "$duration"
        return 0
    else
        echo "-1"  # 失败标记
        return 1
    fi
}

# 测量镜像层信息
measure_image_layers() {
    local image=$1

    # 获取镜像信息
    local image_info=$(ctr image info "$image" 2>/dev/null || echo "{}")
    local layer_count=$(echo "$image_info" | jq -r '.ChainID[]? | length' 2>/dev/null | wc -l)
    local total_size=$(echo "$image_info" | jq -r '.Size' 2>/dev/null || echo "0")

    echo "$layer_count $total_size"
}

# 测量存储使用情况
measure_storage_usage() {
    local snapshotter_path="/var/lib/containerd/io.containerd.snapshotter.v1.$SNAPSHOTTER_NAME"

    if [ -d "$snapshotter_path" ]; then
        local usage=$(du -sb "$snapshotter_path" 2>/dev/null | awk '{print $1}' || echo "0")
        echo "$usage"
    else
        echo "0"
    fi
}

# 性能测试主函数
run_performance_test() {
    log "开始性能测试..."

    # 初始化结果文件
    cat > "$RESULTS_FILE" << EOF
{
    "timestamp": "$(date -Iseconds)",
    "snapshotter": "$SNAPSHOTTER_NAME",
    "test_config": {
        "iterations": $ITERATIONS,
        "warmup_iterations": $WARMUP_ITERATIONS,
        "images_count": ${#TEST_IMAGES[@]}
    },
    "system_info": $(cat "$RESULTS_DIR/system-info-$TIMESTAMP.json"),
    "results": []
}
EOF

    local total_tests=0
    local successful_tests=0

    # 对每个镜像进行测试
    for image in "${TEST_IMAGES[@]}"; do
        info "测试镜像: $image"

        # 获取镜像信息
        local image_info=$(measure_image_layers "$image")
        local layer_count=$(echo "$image_info" | awk '{print $1}')
        local total_size=$(echo "$image_info" | awk '{print $2}')

        info "镜像层数: $layer_count, 总大小: $total_size bytes"

        # 预热运行
        log "预热运行 ($WARMUP_ITERATIONS 次)..."
        for i in $(seq 1 "$WARMUP_ITERATIONS"); do
            measure_startup_time "$image" "warmup-$i" > /dev/null 2>&1 || true
        done

        # 正式测试
        log "正式测试 ($ITERATIONS 次)..."
        local startup_times=()
        local storage_before=$(measure_storage_usage)

        for i in $(seq 1 "$ITERATIONS"); do
            info "第 $i 次测试..."
            local startup_time=$(measure_startup_time "$image" "$i")

            if [ "$startup_time" != "-1" ]; then
                startup_times+=("$startup_time")
                ((successful_tests++))
            else
                warning "第 $i 次测试失败"
            fi

            ((total_tests++))
        done

        # 计算统计数据
        if [ ${#startup_times[@]} -gt 0 ]; then
            local avg_time=$(echo "${startup_times[@]}" | tr ' ' '\n' | awk '{sum+=$1} END {print sum/NR}')
            local min_time=$(echo "${startup_times[@]}" | tr ' ' '\n' | sort -n | head -1)
            local max_time=$(echo "${startup_times[@]}" | tr ' ' '\n' | sort -n | tail -1)

            # 计算标准差
            local variance=$(echo "${startup_times[@]}" | tr ' ' '\n' | awk -v avg="$avg_time" '{sum+=($1-avg)*($1-avg)} END {print sum/NR}')
            local std_dev=$(echo "sqrt($variance)" | bc -l)

            local storage_after=$(measure_storage_usage)
            local storage_delta=$((storage_after - storage_before))

            info "平均启动时间: ${avg_time}s"
            info "最小启动时间: ${min_time}s"
            info "最大启动时间: ${max_time}s"
            info "标准差: ${std_dev}s"
            info "存储使用变化: ${storage_delta} bytes"

            # 更新结果文件
            local temp_file=$(mktemp)
            jq --arg image "$image" \
               --arg layer_count "$layer_count" \
               --arg total_size "$total_size" \
               --arg avg_time "$avg_time" \
               --arg min_time "$min_time" \
               --arg max_time "$max_time" \
               --arg std_dev "$std_dev" \
               --arg storage_delta "$storage_delta" \
               --argjson startup_times "$(printf '%s\n' "${startup_times[@]}" | jq -R . | jq -s .)" \
               '.results += [{
                   "image": $image,
                   "layer_count": ($layer_count | tonumber),
                   "total_size": ($total_size | tonumber),
                   "startup_times": $startup_times,
                   "statistics": {
                       "avg_time": ($avg_time | tonumber),
                       "min_time": ($min_time | tonumber),
                       "max_time": ($max_time | tonumber),
                       "std_dev": ($std_dev | tonumber)
                   },
                   "storage_delta": ($storage_delta | tonumber)
               }]' "$RESULTS_FILE" > "$temp_file"

            mv "$temp_file" "$RESULTS_FILE"
        fi
    done

    # 添加总体统计
    local overall_success_rate=$(echo "scale=2; $successful_tests * 100 / $total_tests" | bc -l)

    local final_temp_file=$(mktemp)
    jq --arg total_tests "$total_tests" \
       --arg successful_tests "$successful_tests" \
       --arg overall_success_rate "$overall_success_rate" \
       '. += {
           "summary": {
               "total_tests": ($total_tests | tonumber),
               "successful_tests": ($successful_tests | tonumber),
               "failed_tests": (($total_tests | tonumber) - ($successful_tests | tonumber)),
               "success_rate": ($overall_success_rate | tonumber)
           }
       }' "$RESULTS_FILE" > "$final_temp_file"

    mv "$final_temp_file" "$RESULTS_FILE"

    success "性能测试完成"
}

# 生成性能报告
generate_performance_report() {
    log "生成性能报告..."

    local report_file="$RESULTS_DIR/performance-report-$TIMESTAMP.txt"

    # 基本报告
    cat > "$report_file" << EOF
=====================================
智能快照器性能测试报告
=====================================
测试时间: $(date)
快照器: $SNAPSHOTTER_NAME
测试镜像数量: ${#TEST_IMAGES[@]}
每个镜像测试次数: $ITERATIONS
预热次数: $WARMUP_ITERATIONS

系统信息:
$(cat "$RESULTS_DIR/system-info-$TIMESTAMP.json" | jq -r '
    "CPU: " + .cpu + "\n" +
    "内存: " + .memory + "\n" +
    "可用磁盘: " + .disk_free + "\n" +
    "内核版本: " + .kernel + "\n" +
    "containerd版本: " + .containerd_version
')

详细结果:
EOF

    # 添加每个镜像的结果
    jq -r '.results[] |
        "\n镜像: " + .image +
        "\n  层数: " + (.layer_count | tostring) +
        "\n  总大小: " + (.total_size | tostring) + " bytes" +
        "\n  平均启动时间: " + (.statistics.avg_time | tostring) + "s" +
        "\n  最小启动时间: " + (.statistics.min_time | tostring) + "s" +
        "\n  最大启动时间: " + (.statistics.max_time | tostring) + "s" +
        "\n  标准差: " + (.statistics.std_dev | tostring) + "s" +
        "\n  存储变化: " + (.storage_delta | tostring) + " bytes\n"
    ' "$RESULTS_FILE" >> "$report_file"

    # 添加总体统计
    jq -r '.summary |
        "\n总体统计:" +
        "\n  总测试数: " + (.total_tests | tostring) +
        "\n  成功数: " + (.successful_tests | tostring) +
        "\n  失败数: " + (.failed_tests | tostring) +
        "\n  成功率: " + (.success_rate | tostring) + "%\n"
    ' "$RESULTS_FILE" >> "$report_file"

    success "性能报告已生成: $report_file"

    # 显示报告摘要
    echo
    echo "性能测试摘要:"
    echo "==============="
    jq -r '.results[] | "\(.image): 平均 \(.statistics.avg_time)s, 最小 \(.statistics.min_time)s, 最大 \(.statistics.max_time)s"' "$RESULTS_FILE"
    echo
    jq -r '.summary | "总成功率: \(.success_rate)% (\(.successful_tests)/\(.total_tests))"' "$RESULTS_FILE"
}

# 主函数
main() {
    log "开始智能快照器性能测试"
    log "====================================="

    # 检查依赖
    check_dependencies

    # 收集系统信息
    collect_system_info

    # 拉取测试镜像
    pull_test_images

    # 运行性能测试
    run_performance_test

    # 生成报告
    generate_performance_report

    log "性能测试完成"
    log "====================================="
    log "结果文件: $RESULTS_FILE"
    log "报告文件: $RESULTS_DIR/performance-report-$TIMESTAMP.txt"
}

# 参数解析
while [[ $# -gt 0 ]]; do
    case $1 in
        --snapshotter)
            SNAPSHOTTER_NAME="$2"
            shift 2
            ;;
        --iterations)
            ITERATIONS="$2"
            shift 2
            ;;
        --warmup)
            WARMUP_ITERATIONS="$2"
            shift 2
            ;;
        --help)
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --snapshotter NAME   快照器名称 (默认: smart)"
            echo "  --iterations NUM     每个镜像的测试次数 (默认: 5)"
            echo "  --warmup NUM         预热次数 (默认: 2)"
            echo "  --help               显示帮助信息"
            exit 0
            ;;
        *)
            error "未知参数: $1"
            exit 1
            ;;
    esac
done

# 运行主函数
main "$@"